package edu.cmu.vlis.distributed.relationalalgebra.visitor;

import java.util.*;

import edu.cmu.vlis.distributed.relationalalgebra.*;

public class PushUnionUpVisitor implements AVisitor<AlgebraNode> {

	@Override
	public AlgebraNode visitNode(Join node) {
		// TODO: finish the function
		// If Union is child of Join push it up.

		node.setLeft(node.getLeft().accept(this));
		node.setRight(node.getRight().accept(this));

		if ((node.getLeft() instanceof Union)
				|| (node.getRight() instanceof Union)) {
			List<AlgebraNode> leftChildren = null;
			List<AlgebraNode> rightChildren = null;
			if (node.getLeft() instanceof Union) {
				leftChildren = ((Union) node.getLeft()).getChildren();
			} else {
				leftChildren = new ArrayList<AlgebraNode>();
				leftChildren.add(node.getLeft());
			}

			if (node.getRight() instanceof Union) {
				rightChildren = ((Union) node.getRight()).getChildren();
			} else {
				rightChildren = new ArrayList<AlgebraNode>();
				rightChildren.add(node.getRight());
			}

			Union u = new Union();
			List<AlgebraNode> newChildren = new ArrayList<AlgebraNode>();
			for (AlgebraNode leftChild : leftChildren) {
				for (AlgebraNode rightChild : rightChildren) {
					Join newJoin = new Join(leftChild, node.getLeftColumn(),
							rightChild, node.getRightColumn());
					newJoin.accept(this);
					newChildren.add(newJoin);
				}
			}
			u.setChildren(newChildren);
			return u;
		}
		
		return node;
	}

	@Override
	public AlgebraNode visitNode(Selection node) {
		node.setChild(node.getChild().accept(this));
		if (node.getChild() instanceof Union) {
			Union u = (Union) node.getChild();
			List<AlgebraNode> newChildren = new ArrayList<AlgebraNode>();
			for (AlgebraNode child : u.getChildren()) {
				Selection newChild = new Selection(child, node.getCondition());
				newChildren.add(newChild);
			}
			u.setChildren(newChildren);
			return u;
		} else {
			return node;
		}
	}

	@Override
	public AlgebraNode visitNode(Relation node) {
		return node;
	}

	@Override
	public AlgebraNode visitNode(Union node) {
		List<AlgebraNode> newChildren = new ArrayList<AlgebraNode>();

		for (AlgebraNode child : node.getChildren()) {
			AlgebraNode newChild = child.accept(this);
			if (newChild instanceof Union) {
				newChildren.addAll(((Union) newChild).getChildren());
			} else {
				newChildren.add(newChild);
			}
		}
		node.setChildren(newChildren);
		return node;
	}

	@Override
	public AlgebraNode visitNode(Fragment node) {
		return node;
	}

	@Override
	public AlgebraNode visitNode(Projection node) {
		node.setChild(node.getChild().accept(this));
		if (node.getChild() instanceof Union) {
			Union u = (Union) node.getChild();
			List<AlgebraNode> newChildren = new ArrayList<AlgebraNode>();
			for (AlgebraNode child : u.getChildren()) {
				Projection newChild = new Projection(child, node.getColumns());
				newChildren.add(newChild);
			}
			u.setChildren(newChildren);
			return u;
		} else {
			return node;
		}
	}

}
